#!/bin/bash

set -e

# 配置参数
RUNS=5  # 测试次数
IMAGE_ESGZ="randkin/llm-environment:esgz"
IMAGE_ESGZV2="randkin/llm-environment:esgzv2"
RESULT_DIR="../result"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)

# 创建结果目录
mkdir -p "$RESULT_DIR"
CSV_FILE="$RESULT_DIR/benchmark_$TIMESTAMP.csv"
SUMMARY_FILE="$RESULT_DIR/summary_$TIMESTAMP.txt"
PLOT_FILE="$RESULT_DIR/plot_$TIMESTAMP.png"

# 存储结果的时间数组
declare -a esgz_total_times
declare -a esgzv2_total_times

echo "========================================="
echo "  大模型推理容器冷启动对比测试"
echo "  测试次数: $RUNS"
echo "  原生方案镜像: $IMAGE_ESGZ"
echo "  优化方案镜像: $IMAGE_ESGZV2"
echo "  结果目录: $(realpath "$RESULT_DIR")"
echo "========================================="

# 创建CSV文件头
echo "TestRun,Image,PullTime(s),ColdStartTime(s),TotalTime(s)" > "$CSV_FILE"

# 测试函数
run_test() {
    local image="$1"
    local run_num="$2"
    
    # 清理环境
    sudo nerdctl rmi -f "$image" 2>/dev/null || true
    sudo systemctl restart containerd
    sleep 2  # 等待服务稳定
    
    # 运行测试并捕获输出
    output=$(./optimized_test.sh "$image")
    echo "$output"
    
    # 提取时间数据
    pull_time=$(echo "$output" | grep "镜像拉取耗时" | awk '{print $4}')
    cold_start_time=$(echo "$output" | grep "【容器冷启动时间】" | awk '{print $2}')
    total_time=$(echo "$output" | grep "【镜像拉取+冷启动时间】" | awk '{print $2}')
    
    # 记录到CSV
    echo "$run_num,$image,$pull_time,$cold_start_time,$total_time" >> "$CSV_FILE"
    
    # 返回总时间
    echo "$total_time"
}

# 测试 esgz 镜像
echo "开始测试原生方案 ($IMAGE_ESGZ)..."
for ((i=1; i<=$RUNS; i++)); do
    echo "----- 第 $i 次测试 -----"
    total_time=$(run_test "$IMAGE_ESGZ" "$i")
    esgz_total_times+=("$total_time")
    echo "记录总时间: $total_time 秒"
done

# 测试 esgzv2 镜像
echo "开始测试优化方案 ($IMAGE_ESGZV2)..."
for ((i=1; i<=$RUNS; i++)); do
    echo "----- 第 $i 次测试 -----"
    total_time=$(run_test "$IMAGE_ESGZV2" "$i")
    esgzv2_total_times+=("$total_time")
    echo "记录总时间: $total_time 秒"
done

# 计算结果统计
function calculate_stats {
    local times=("$@")
    local sum=0
    local min=999999
    local max=0
    
    for t in "${times[@]}"; do
        sum=$(echo "$sum + $t" | bc)
        if (( $(echo "$t < $min" | bc -l) )); then min=$t; fi
        if (( $(echo "$t > $max" | bc -l) )); then max=$t; fi
    done
    
    local avg=$(echo "scale=3; $sum / ${#times[@]}" | bc)
    echo "$avg $min $max"
}

# 计算统计信息
esgz_stats=($(calculate_stats "${esgz_total_times[@]}"))
esgzv2_stats=($(calculate_stats "${esgzv2_total_times[@]}"))

# 计算提升百分比
improvement=$(echo "scale=2; (${esgz_stats[0]} - ${esgzv2_stats[0]}) / ${esgz_stats[0]} * 100" | bc)

# 生成汇总报告
{
echo "========================================="
echo "             测试结果汇总"
echo "========================================="
printf "%-15s %-10s %-10s %-10s\n" "方案" "平均(秒)" "最小(秒)" "最大(秒)"
printf "%-15s %-10.3f %-10.3f %-10.3f\n" "原生(esgz)" ${esgz_stats[0]} ${esgz_stats[1]} ${esgz_stats[2]}
printf "%-15s %-10.3f %-10.3f %-10.3f\n" "优化(esgzv2)" ${esgzv2_stats[0]} ${esgzv2_stats[1]} ${esgzv2_stats[2]}
echo "-----------------------------------------"
echo "总启动时间提升: $improvement%"
echo "-----------------------------------------"

# 检查是否达到20%提升目标
if (( $(echo "$improvement >= 20" | bc -l) )); then
    echo "✅ 达到目标: 总启动时间降低超过20%"
else
    echo "❌ 未达目标: 需要至少20%的提升"
fi
echo "========================================="
} | tee "$SUMMARY_FILE"

# 生成图表
echo "生成性能对比图表..."
python3 - <<EOF
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import os

# 读取CSV数据
df = pd.read_csv("$CSV_FILE")

# 设置图表
plt.figure(figsize=(12, 8))
plt.style.use('seaborn-darkgrid')

# 为每个镜像创建单独的数据框
df_esgz = df[df['Image'] == '$IMAGE_ESGZ']
df_esgzv2 = df[df['Image'] == '$IMAGE_ESGZV2']

# 创建分组柱状图
bar_width = 0.35
index = np.arange(len(df_esgz))

plt.bar(index, df_esgz['TotalTime(s)'], bar_width, label='原生方案 ($IMAGE_ESGZ)', color='#1f77b4')
plt.bar(index + bar_width, df_esgzv2['TotalTime(s)'], bar_width, label='优化方案 ($IMAGE_ESGZV2)', color='#ff7f0e')

# 添加数据标签
for i, val in enumerate(df_esgz['TotalTime(s)']):
    plt.text(i, val + 0.05, f'{val:.1f}', ha='center', va='bottom')
for i, val in enumerate(df_esgzv2['TotalTime(s)']):
    plt.text(i + bar_width, val + 0.05, f'{val:.1f}', ha='center', va='bottom')

# 设置图表标题和标签
plt.title('容器启动时间对比 (总时间 = 拉取时间 + 冷启动时间)', fontsize=16)
plt.xlabel('测试轮次', fontsize=14)
plt.ylabel('总时间 (秒)', fontsize=14)
plt.xticks(index + bar_width/2, [f'第 {i+1} 次' for i in range(len(df_esgz))])
plt.legend(fontsize=12)

# 添加网格线
plt.grid(True, linestyle='--', alpha=0.7)

# 添加性能提升标注
avg_esgz = ${esgz_stats[0]}
avg_esgzv2 = ${esgzv2_stats[0]}
improvement = $improvement
plt.text(len(df_esgz)/2 - 0.5, max(df_esgz['TotalTime(s)']) * 0.9, 
         f'平均提升: {improvement}%', 
         fontsize=14, ha='center', bbox=dict(facecolor='yellow', alpha=0.5))

# 保存图表
plt.tight_layout()
plt.savefig("$PLOT_FILE", dpi=300)
EOF

echo "========================================="
echo "测试完成!"
echo "CSV数据文件: $(realpath "$CSV_FILE")"
echo "汇总报告: $(realpath "$SUMMARY_FILE")"
echo "性能图表: $(realpath "$PLOT_FILE")"
echo "========================================="